package com.frank.oj.core.dispatcher;

import cn.hutool.core.lang.UUID;
import cn.hutool.json.JSONObject;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.frank.model.dto.CompileDTO;
import com.frank.model.dto.TestJudgeReq;
import com.frank.model.dto.TestJudgeRes;
import com.frank.model.dto.ToJudgeDTO;
import com.frank.model.entity.judge.Judge;
import com.frank.model.entity.judge.JudgeServer;
import com.frank.model.entity.judge.RemoteJudgeAccount;
import com.frank.oj.common.result.CommonResult;
import com.frank.oj.common.result.ResultStatus;
import com.frank.oj.constant.Constants;
import com.frank.oj.core.ChooseUtils;
import com.frank.oj.service.entity.judge.JudgeEntityService;
import com.frank.oj.service.entity.judge.JudgeServerEntityService;
import com.frank.oj.service.entity.judge.RemoteJudgeAccountEntityService;
import com.frank.oj.utils.RedisUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestTemplate;

import javax.annotation.Resource;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 分发到判题机
 *
 * @author frank
 * @Date 2024/4/24
 */
@Component
@Slf4j(topic = "hcoj")
public class Dispatcher {
    @Resource
    private RestTemplate restTemplate;

    @Resource
    private JudgeServerEntityService judgeServerEntityService;

    @Resource
    private JudgeEntityService judgeEntityService;

    @Resource
    private ChooseUtils chooseUtils;

    @Resource
    private RedisUtils redisUtils;
    @Resource
    private RemoteJudgeAccountEntityService remoteJudgeAccountEntityService;

    private final static ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(20);

    private final static Map<String, Future> futureTaskMap = new ConcurrentHashMap<>(20);

    // 每个提交任务尝试300次失败则判为提交失败
    protected final static Integer maxTryNum = 300;

    public CommonResult<Object> dispatch(Constants.TaskType taskType, Object data) {
        switch (taskType) {
            case JUDGE:
                defaultJudge((ToJudgeDTO) data, taskType.getPath());
                break;
            case REMOTE_JUDGE:
                remoteJudge((ToJudgeDTO) data, taskType.getPath());
                break;
            case TEST_JUDGE:
                testJudge((TestJudgeReq) data, taskType.getPath());
                break;
            case COMPILE_SPJ:
            case COMPILE_INTERACTIVE:
                return compile((CompileDTO) data, taskType.getPath());
            default:
                throw new IllegalArgumentException("判题机不支持此调用类型");
        }
        return null;
    }

    /**
     * 测试判题
     */
    private void testJudge(TestJudgeReq testJudgeReq, String path) {

        AtomicInteger count = new AtomicInteger(0);
        String taskKey = testJudgeReq.getUniqueKey();
        Runnable getResultTask = () -> {
            if (count.get() > maxTryNum) {
                releaseTaskThread(taskKey);
                return;
            }
            count.getAndIncrement();
            JudgeServer judgeServer = chooseUtils.chooseServer(false);
            if (judgeServer != null) {
                try {
                    String url = "http://" + judgeServer.getUrl() + path;
                    JSONObject resultJson = restTemplate.postForObject(url, testJudgeReq, JSONObject.class);

                    if (resultJson != null) {
                        if (resultJson.getInt("status") == ResultStatus.SUCCESS.getStatus()) {
                            TestJudgeRes testJudgeRes = resultJson.getBean("data", TestJudgeRes.class);
                            testJudgeRes.setInput(testJudgeReq.getTestCaseInput());
                            testJudgeRes.setExpectedOutput(testJudgeReq.getExpectedOutput());
                            testJudgeRes.setProblemJudgeMode(testJudgeReq.getProblemJudgeMode());

                            redisUtils.set(testJudgeReq.getUniqueKey(), testJudgeRes, 60);
                        } else {
                            TestJudgeRes testJudgeRes = TestJudgeRes.builder()
                                    .status(Constants.Judge.STATUS_SYSTEM_ERROR.getStatus())
                                    .time(0L)
                                    .memory(0L)
                                    .stderr(resultJson.getStr("msg"))
                                    .build();
                            redisUtils.set(testJudgeReq.getUniqueKey(), testJudgeRes, 60);
                        }
                    }

                } catch (Exception e) {
                    log.error("[Test Judge] Request the judge server [" + judgeServer.getUrl() + "] error-------------->", e);
                    TestJudgeRes testJudgeRes = TestJudgeRes.builder()
                            .status(Constants.Judge.STATUS_SYSTEM_ERROR.getStatus())
                            .time(0L)
                            .memory(0L)
                            .stderr("Failed to connect the judgeServer. Please resubmit this submission again!")
                            .build();
                    redisUtils.set(testJudgeReq.getUniqueKey(), testJudgeRes, 60);
                } finally {
                    releaseJudgeServer(judgeServer.getId());
                    releaseTaskThread(taskKey);
                }

            }
        };

        ScheduledFuture<?> scheduledFuture = scheduler.scheduleWithFixedDelay(getResultTask, 0, 1, TimeUnit.SECONDS);
        futureTaskMap.put(taskKey, scheduledFuture);
    }

    /**
     * 默认判题
     */
    private void defaultJudge(ToJudgeDTO data, String path) {
        Long submitId = data.getJudge().getSubmitId();
        AtomicInteger count = new AtomicInteger(0);
        String taskKey = UUID.randomUUID().toString() + submitId;

        Runnable getResultTask = () -> {
            if (count.get() > maxTryNum) {
                checkResult(null, submitId);
                releaseTaskThread(taskKey);
                return;
            }

            count.getAndIncrement();
            JudgeServer judgeServer = chooseUtils.chooseServer(false);
            if (judgeServer != null) {
                CommonResult result = null;
                try {
                    result = restTemplate.postForObject("http://" + judgeServer.getUrl() + path, data, CommonResult.class);
                } catch (RestClientException e) {
                    log.error("[Self Judge] Request the judge server [" + judgeServer.getUrl() + "] error -------------->", e);
                } finally {
                    checkResult(result, submitId);
                    releaseJudgeServer(judgeServer.getId());
                    releaseTaskThread(taskKey);
                }

            }
        };
        ScheduledFuture<?> scheduledFuture = scheduler.scheduleWithFixedDelay(getResultTask, 0, 2, TimeUnit.SECONDS);
        futureTaskMap.put(taskKey, scheduledFuture);
    }


    /**
     * 远程判题
     */
    private void remoteJudge(ToJudgeDTO data, String path) {
        String oj = data.getRemoteJudgeProblem().split("-")[0];
        if (oj.equals("GYM")) {
            oj = "CF";
        }

        // 如果是vj判題，同时不要是以已经有提交id的获取结果操作，归属于CF的判題，需控制判题机的权限，一机一题
        boolean isCFFixServerJudge = ChooseUtils.openCodeforcesFixServer
                && !data.getIsHasSubmitIdRemoteReJudge()
                && Constants.RemoteOJ.CODEFORCES.getName().equals(oj);

        Long submitId = data.getJudge().getSubmitId();
        AtomicInteger count = new AtomicInteger(0);
        String taskKey = UUID.randomUUID().toString() + submitId;
        final String finalOj = oj;

        Runnable getResultTask = () -> {
            if (count.get() > maxTryNum) { // 300次失败则判为提交失败
                // 远程判题需要将账号归为可用
                changeRemoteJudgeStatus(finalOj, data.getUsername(), null);
                checkResult(null, submitId);
                releaseTaskThread(taskKey);
                return;
            }
            count.getAndIncrement();
            JudgeServer judgeServer = null;

            if (!isCFFixServerJudge) {
                judgeServer = chooseUtils.chooseServer(true);
            } else {
                judgeServer = chooseUtils.chooseFixedServer(true, "cf_submittable", data.getIndex(), data.getSize());
            }

            if (judgeServer != null) { // 获取到判题机资源
                data.setJudgeServerIp(judgeServer.getIp());
                data.setJudgeServerPort(judgeServer.getPort());

                CommonResult result = null;
                try {
                    result = restTemplate.postForObject("http://" + judgeServer.getUrl() + path, data, CommonResult.class);
                } catch (Exception e) {
                    log.error("[Remote Judge] Request the judge server [" + judgeServer.getUrl() + "] error-------------->", e);
                    changeRemoteJudgeStatus(finalOj, data.getUsername(), judgeServer);
                } finally {
                    checkResult(result, submitId);
                    if (!isCFFixServerJudge) {
                        // 无论成功与否，都要将对应的当前判题机当前判题数减1
                        releaseJudgeServer(judgeServer.getId());
                    }
                    releaseTaskThread(taskKey);
                }
            }
        };

        // 重试
        ScheduledFuture<?> scheduledFuture = scheduler.scheduleWithFixedDelay(getResultTask, 0, 2, TimeUnit.SECONDS);
        futureTaskMap.put(taskKey, scheduledFuture);
    }

    private void changeRemoteJudgeStatus(String oj, String username, JudgeServer judgeServer) {
        changeAccountStatus(oj, username);
        if (ChooseUtils.openCodeforcesFixServer) {
            if (oj.equals(Constants.RemoteOJ.CODEFORCES.getName())
                    || oj.equals(Constants.RemoteOJ.GYM.getName())) {
                if (judgeServer != null) {
                    changeServerSubmitCFStatus(judgeServer.getIp(), judgeServer.getPort());
                }
            }
        }
    }

    private void changeServerSubmitCFStatus(String ip, Integer port) {
        if (StringUtils.isEmpty(ip) || port == null) {
            return;
        }
        LambdaUpdateWrapper<JudgeServer> judgeServerUpdateWrapper = new LambdaUpdateWrapper<>();
        judgeServerUpdateWrapper.set(JudgeServer::getCfSubmittable, true)
                .eq(JudgeServer::getIp, ip)
                .eq(JudgeServer::getIsRemote, true)
                .eq(JudgeServer::getPort, port);
        boolean isOk = judgeServerEntityService.update(judgeServerUpdateWrapper);

        if (!isOk) { // 重试8次
            tryAgainUpdateServer(judgeServerUpdateWrapper, ip, port);
        }
    }

    private void tryAgainUpdateServer(LambdaUpdateWrapper<JudgeServer> updateWrapper, String ip, Integer port) {
        boolean retryable;
        int attemptNumber = 0;
        do {
            boolean success = judgeServerEntityService.update(updateWrapper);
            if (success) {
                return;
            } else {
                attemptNumber++;
                retryable = attemptNumber < 8;
                if (attemptNumber == 8) {
                    log.error("[Remote Judge] Change Codeforces Judge Server Status to `true` Failed! =======>{}", "ip:" + ip + ",port:" + port);
                    break;
                }
                try {
                    Thread.sleep(300);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } while (retryable);
    }

    private void changeAccountStatus(String oj, String username) {
        LambdaUpdateWrapper<RemoteJudgeAccount> updateWrapper = new LambdaUpdateWrapper<>();
        updateWrapper.set(RemoteJudgeAccount::getStatus, true)
                .eq(RemoteJudgeAccount::getStatus, false)
                .eq(RemoteJudgeAccount::getUsername, username);

        if (oj.equals("GYM")) {
            oj = "CF";
        }
        updateWrapper.eq(RemoteJudgeAccount::getOj, oj);

        boolean updated = remoteJudgeAccountEntityService.update(updateWrapper);
        if (!updated) {
            tryAgainUpdateAccount(updateWrapper, oj, username);
        }
    }

    private void tryAgainUpdateAccount(LambdaUpdateWrapper<RemoteJudgeAccount> updateWrapper, String oj, String username) {
        boolean retryable;
        int attemptNumber = 0;
        do {
            boolean success = remoteJudgeAccountEntityService.update(updateWrapper);
            if (success) {
                return;
            } else {
                attemptNumber++;
                retryable = attemptNumber < 8;
                if (attemptNumber == 8) {
                    log.error("[Remote Judge] Failed to change the account to be available----------->{}", "oj:" + oj + ",username:" + username);
                    break;
                }
                try {
                    Thread.sleep(300);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } while (retryable);
    }

    /**
     * 结果处理
     */
    private void checkResult(CommonResult<Void> result, Long submitId) {

        Judge judge = new Judge();
        if (result == null) {
            judge.setSubmitId(submitId);
            judge.setStatus(Constants.Judge.STATUS_SUBMITTED_FAILED.getStatus());
            judge.setErrorMessage("Failed to connect the judgeServer. Please resubmit this submission again!");
            judgeEntityService.updateById(judge);
        } else {
            if (result.getStatus() != ResultStatus.SUCCESS.getStatus()) { // 如果是结果码不是200 说明调用有错误
                // 判为系统错误
                judge.setStatus(Constants.Judge.STATUS_SYSTEM_ERROR.getStatus())
                        .setErrorMessage(result.getMsg());
                judgeEntityService.updateById(judge);
            }
        }
    }

    /**
     * 编译特殊判题程序或交互程序
     */
    private CommonResult compile(CompileDTO data, String path) {
        CommonResult result = CommonResult.errorResponse("没有可用的判题服务器，请重新尝试！");
        JudgeServer judgeServer = chooseUtils.chooseServer(false);
        if (judgeServer != null) {
            try {
                result = restTemplate.postForObject("http://" + judgeServer.getUrl() + path, data, CommonResult.class);
            } catch (Exception e) {
                log.error("[Compile] Request the judge server [" + judgeServer.getUrl() + "] error-------------->", e.getMessage());
            } finally {
                // 无论成功与否，都要将对应的当前判题机当前判题数减1
                releaseJudgeServer(judgeServer.getId());
            }
        }
        return result;
    }

    /**
     * 释放指定任务的线程。
     * 该方法会尝试取消与给定任务键相关联的未来任务，并从未来任务映射中移除该任务。
     *
     * @param taskKey 任务的键，用于在futureTaskMap中查找对应的Future任务。
     */
    private void releaseTaskThread(String taskKey) {
        // 根据任务键获取Future对象
        Future future = futureTaskMap.get(taskKey);
        if (future != null) {
            // 尝试取消任务，并检查是否成功取消
            boolean cancelled = future.cancel(true);
            if (cancelled) {
                // 如果任务成功取消，则从映射中移除该任务
                futureTaskMap.remove(taskKey);
            }
        }
    }

    /**
     * 释放当前判题机
     */
    private void releaseJudgeServer(Integer judgeServerId) {
        LambdaUpdateWrapper<JudgeServer> updateWrapper = new LambdaUpdateWrapper<>();
        updateWrapper.setSql("task_number = task_number-1").eq(JudgeServer::getId, judgeServerId);

        boolean updated = judgeServerEntityService.update(updateWrapper);
        if (!updated) {
            tryAgainUpdateJudge(updateWrapper);
        }
    }

    /**
     * 尝试重新更新判题机
     */
    private void tryAgainUpdateJudge(LambdaUpdateWrapper<JudgeServer> updateWrapper) {
        boolean retryable;
        int attemptNumber = 0;
        do {
            boolean success = judgeServerEntityService.update(updateWrapper);
            if (success) {
                return;
            } else {
                attemptNumber++;
                retryable = attemptNumber < 8;
                if (attemptNumber == 8) {
                    break;
                }
                try {
                    Thread.sleep(300);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } while (retryable);
    }
}
